import pymongo as pm
import json
import torch
from PyCmpltrtok.util_mongo import mongo_upsert, mongo_get, VALUE
from PyCmpltrtok.common import md5, sep

VER_INT = 1
VERSION_COL = 'version'
EMBED_COL = 'embed'
USERNAME = 'embed'
EMBED_TBL = 'sentence_bank'


def sentence2embed(dev, model, tokenizer, xsent, max_seq_length=512):
    ids = tokenizer(
        xsent,
        max_length=max_seq_length,
        truncation=True,
        padding=False,
        # padding='max_length',  # https://stackoverflow.com/questions/70067608/how-padding-in-huggingface-tokenizer-works
        # https://stackoverflow.com/questions/61443480/huggingfaces-bert-tokenizer-not-adding-pad-token
        return_tensors='pt',
    )
    with torch.no_grad():
        xout = model(**{k: v.to(dev) for k, v in ids.items()})
        embed = xout['pooler_output']
    return embed.tolist()[0]


def get_each_line_of_file(xpath, limit=0, enc='utf8'):
    """
    返回可获取文件的每一行的generator。

    :param xpath: 文件路径
    :param limit: 限制行数，0为不限制，默认0。
    :param enc: 编码，默认utf8
    :return: 可获取文件的每一行的generator
    """
    cnt = -1
    with open(xpath, 'r', encoding=enc) as fi:
        while True:
            # 读取一行
            xline = fi.readline()
            if not xline:
                break
            if '\r\n' == xline[-2:]:
                xline = xline[:-2]
            elif '\r' == xline[-1:]:
                xline = xline[:-1]
            xline = xline.strip()
            if not xline:
                continue

            # 人为限定读取的行数
            cnt += 1
            if limit and cnt >= limit:
                break

            yield xline


def get_cblue_text_from_json_text(xline):
    """
    按CBLUE CMeIE v2数据集的格式返回文本。

    :param xline: 文本（JSON格式）
    :return: JSON文本解析后的['text']值，或不能解析等其他情况返回None。
    """
    xtext = None
    try:
        xdict = json.loads(xline)
        xtext = xdict['text']
    except Exception as ex:
        print(ex)
    return xtext


if '__main__' == __name__:

    from transformers import AutoModel, AutoTokenizer
    from transformers import AutoConfig

    def _main():
        # 连接Mongodb
        mongo = pm.MongoClient('127.0.0.1', 27017, serverSelectionTimeoutMS=3000)
        mdb = mongo['CBLUE']

        # 加载文本BERT
        dev = torch.device(0)
        if 0:
            model_name = 'bert-base-chinese'  # 用模型名，需联网，需翻墙
        else:
            model_name = r'C:\Users\peter\.cache\huggingface\hub\models--bert-base-chinese\snapshots\8d2a91f91cc38c96bb8b4556ba70c392f8d5ee55'  # 使用下载后的路径，不需联网
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = AutoModel.from_pretrained(model_name).to(dev)

        # 确定最大token长度
        xlen01 = tokenizer.model_max_length
        config = AutoConfig.from_pretrained(model_name)
        xlen02 = config.max_position_embeddings
        max_len = min(xlen01, xlen02)

        xpath_list = [
            r'D:\_dell7590_root\local\LNP_datasets\med\CBLUE\CMeIE-V2\CMeIE-V2_train.jsonl',
        ]

        for offset, xpath in enumerate(xpath_list):
            sep(xpath)
            xgen = get_each_line_of_file(xpath, 1000)
            for i, xline in enumerate(xgen):
                print('.', end='')
                if i % 50 == 0:
                    print(offset, i)

                # print(i, xline)
                xtext = get_cblue_text_from_json_text(xline)
                xmd5 = md5(xtext)
                # print(i, xtext, xmd5)

                # 获取以有数据
                xdict = mongo_get(mdb, EMBED_TBL, USERNAME, xmd5, only_value=False)

                # 已经存在该句子的嵌入
                if xdict and xdict.get(EMBED_COL, None):
                    # print(f'Skip {xmd5} {xtext}')
                    continue

                # 做嵌入
                xembed = sentence2embed(dev, model, tokenizer, xtext, max_seq_length=max_len)
                print('+', end='')

                # 入库
                mongo_upsert(mdb, EMBED_TBL, USERNAME, xmd5, {
                    VERSION_COL: VER_INT,
                    VALUE: xtext,
                    EMBED_COL: xembed
                })
                # print(f'Embed {xmd5} {xtext}')
            print()
            sep(f'{xpath} over')

    _main()
